import json
import numpy as np
from StringIO import StringIO

#convert an array to a string
def array_2_string(array):
    str_out = StringIO()
    np.savetxt(str_out, array)
    return str_out.getvalue()

# convert a string to an array
def string_2_array(string):
    str_in = StringIO(string)
    return np.loadtxt(str_in)

def main():
    pdnn_file='exp_pdnn/warm6_str_dnn_fmllr/dnn.temp'
    #exp/dnn6_fbank_dbn_dnn/final_txt.nnet
    #exp/dnn6_fbank_dbn/6_txt.dbn
    kaldi_file=open('exp/dnn6_fmllr_1124-dbn_dnn/final_txt.nnet','r')


    num_hidden=1124
    num_hlayers=6
    num_output=2042

    count=0
    current_layer=-1
    read_weigts=False
    read_bias=False

    weight_array=np.zeros(num_hidden)

    nnet_dict={}

    for line in kaldi_file:
        line=line.rstrip()
        if line.startswith('<LearnRateCoef'):
            current_layer=current_layer+1
            read_weigts=True
            read_bias=False
            count=1
        elif (current_layer==num_hlayers) and read_weigts:
            key = 'logreg W'
            temp_array=string_2_array((line.strip()).replace(']',''))
   
            if count==1:
                print key
                weight_array = temp_array
                #print temp_array.size
            else:
                weight_array = np.vstack((weight_array,temp_array))
            
       
            if count==num_output:
                read_bias=True
                read_weigts=False
                #print weight_array.shape
                nnet_dict[key]=array_2_string(np.transpose(weight_array))
            count=count+1
        
        elif read_weigts:
            key = str(current_layer) + ' sigmoid W'
            temp_array=string_2_array((line.strip()).replace(']',''))
            if count==1:
                print key
                weight_array = temp_array
                #print temp_array.size
            else:
                weight_array = np.vstack((weight_array,temp_array))
        
            if count==num_hidden:
                read_bias=True
                read_weigts=False
                #print weight_array.shape
                nnet_dict[key]= array_2_string(np.transpose(weight_array))
            count=count+1
        elif (current_layer==num_hlayers) and read_bias:
            key = 'logreg b'
            print key
            read_bias=False
            temp=line.split()
            value = '\n'.join(str(p) for p in temp[1:num_output+1])+'\n'
            nnet_dict[key]=value
        elif read_bias:
            key = str(current_layer) + ' sigmoid b'
            print key
            read_bias=False
            temp=line.split()
            value = '\n'.join(str(p) for p in temp[1:num_hidden+1])+'\n'
            nnet_dict[key]=value
        else:
            print line
    print len(nnet_dict)
    with open(pdnn_file, 'wb') as fp:
            json.dump(nnet_dict, fp, indent=2, sort_keys = True)
            fp.flush()

if __name__=="__main__":
    main()

